package jtg.generator;

import com.microsoft.z3.*;
import javafx.util.Pair;
import jtg.graphics.SootCFG;
import jtg.parser.Z3Parser;
import jtg.state.BranchPathState;
import jtg.state.PathState;
import jtg.visualizer.Visualizer;
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.Value;
import soot.jimple.internal.*;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class BranchCoverageGenerator extends Generator {

    Set<Pair<Unit,Unit>> visitedIfNodeNext = new HashSet<>();

    Set<Unit> allVisitedNode = new HashSet<>();

    int branchNum = 0;

    public BranchCoverageGenerator(String className, String methodName) {
        super(className, methodName);
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new PrimeCoverageGenerator(defaultClsPath, className, methodName);
    }

    public BranchCoverageGenerator(String classPath, String className, String methodName) {
        super(classPath, className, methodName);
    }

    @Override
    void drawCFG() {
        drawCFG(clsName + "." + mtdName, false);
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }

    private List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }

    private List<Local> getParameters() {
        return new ArrayList<>(body.getParameterLocals());
    }

    public List<String> generate() {

        ArrayList<String> testSet = null;

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            Context context = new Context();
            Solver solver = context.mkSolver();

            if(ug.getHeads().size()>1) throw new Exception("The number of heads is more than 1");
            Unit h = ug.getHeads().get(0);
            // path = ug.getExtendedBasicBlockPathBetween(h, t);
            for (Unit unit: ug) {
                if(unit instanceof JIfStmt){
                    branchNum++;
                }
            }
            List<BranchPathState> pathStates = initPathState(context, h);
            calculateWithBFS(context, ug, new HashSet<Unit>(), pathStates, testSet);
            int allBranchNum = branchNum*2;
            int visitedBranchNum = visitedIfNodeNext.size();
            System.out.println("The number of all branches is: " + allBranchNum + ", the number of visited branches is: " + visitedBranchNum);


        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
            e.printStackTrace();
        }
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc.toString());
                count++;
            }
        }
        return testSet;
    }

    public List<BranchPathState> initPathState(Context context, Unit h){
        List<BranchPathState> pathStates = new ArrayList<>();
        // 创建初始状态
        BranchPathState initialState = new BranchPathState(h);
        for(Local local : getParameters()){
            String parName = local.getName();
            initialState.params.put(parName, context.mkIntConst(parName));
            initialState.paramLatestMap.put(parName,0);
        }
        initialState.unitPath.add(h);
        pathStates.add(initialState);
        return pathStates;
    }

    public void calculateWithBFS(Context context, UnitGraph ug, Set<Unit> visitedNode, List<BranchPathState> pathStates, List<String> testSet) throws Exception {
        while(!pathStates.isEmpty()) {
            BranchPathState pathState = pathStates.remove(0);
            Unit unit = pathState.node;
            if (pathState.unitPath.size() > 100 || unit instanceof JReturnStmt || unit instanceof JReturnVoidStmt) {
                if (pathCheck(context, pathState)) {
                    List<BoolExpr> pathConstraint = calPathConstraint(context, pathState.unitPath);
                    String testSetByPathConstraint = getTestSetByPathConstraint(context, pathConstraint);
                    if (testSetByPathConstraint != null) {
                        visitedIfNodeNext.addAll(pathState.branchRecord);
                        testSet.add(testSetByPathConstraint);
                    }
                }
                continue;
            }

            Solver solver = context.mkSolver();
            List<BoolExpr> pathConstraint1 = calPathConstraint(context, pathState.unitPath);
            if(!pathConstraint1.isEmpty()) {
                for(BoolExpr expr : pathConstraint1 ){
                    solver.add(expr);
                }

                if(solver.check() == Status.UNSATISFIABLE){
                    continue;
                }
            }

            visitedNode.add(unit);
            if (unit instanceof JIfStmt) {
                // 如果是if语句，那么就是两条路径
                // 生成两个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                BranchPathState pathState2 = new BranchPathState(pathState);
                // 生成两个新的约束
                // 有可能是一个循环，判断一下

                // BoolExpr ifstms = (BoolExpr) getZ3ExprFromSootExpr(context, pathState.params, ((JIfStmt) unit).getCondition());
                List<Unit> nextUnitList = ug.getSuccsOf(unit);
                allVisitedNode.add(unit);
                Unit nextUnit1 = nextUnitList.get(0);
                Unit nextUnit2 = nextUnitList.get(1);
                pathState1.node = nextUnit1;
                pathState2.node = nextUnit2;
                pathState1.branchRecord.add(new Pair<>(unit,nextUnit1));
                pathState2.branchRecord.add(new Pair<>(unit,nextUnit2));
                pathState1.unitPath.add(nextUnit1);
                pathState2.unitPath.add(nextUnit2);
                // pathState1.PathConstraints.add(context.mkNot(ifstms));
                // pathState2.PathConstraints.add(ifstms);


                pathStates.add(pathState1);
                pathStates.add(pathState2);
            } else if (unit instanceof JAssignStmt) {
                // 如果是赋值语句，那么就是一条路径
                // 生成一个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                // 生成一个新的约束
                String leftOp = ((JAssignStmt) unit).getLeftOp().toString();
                pathState1.params.put(leftOp, context.mkIntConst(leftOp));
                // Expr rightOp = getZ3ExprFromSootExpr(context, pathState1.params, ((JAssignStmt) unit).getRightOp());
                // if (rightOp == null) {
                //    System.out.println("Ignore the statement: " + unit.toString());
                //    continue;
                //}
                // pathState1.PathConstraints.add(context.mkEq(pathState1.params.get(leftOp), rightOp));
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);

            } else if (unit instanceof JGotoStmt) {
                // 如果是goto语句，那么就是一条路径
                // 生成一个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                // 判断是否是循环
                Unit nextUnit = ug.getSuccsOf(unit).get(0);

                pathState1.node = nextUnit;
                pathState1.unitPath.add(pathState1.node);

                pathStates.add(pathState1);
            } else if(unit instanceof JLookupSwitchStmt){
                // 如果是switch语句，那么就是多条路径
                // 生成多个新的状态
                List<Unit> nextUnitList = ug.getSuccsOf(unit);
                allVisitedNode.add(unit);
                for(Unit nextUnit : nextUnitList){
                    BranchPathState pathState1 = new BranchPathState(pathState);
                    pathState1.node = nextUnit;
                    pathState1.unitPath.add(pathState1.node);
                    pathStates.add(pathState1);
                }
            } else if(unit instanceof JTableSwitchStmt){
                // 如果是switch语句，那么就是多条路径
                // 生成多个新的状态
                List<Unit> nextUnitList = ug.getSuccsOf(unit);
                allVisitedNode.add(unit);
                for(Unit nextUnit : nextUnitList){
                    BranchPathState pathState1 = new BranchPathState(pathState);
                    pathState1.node = nextUnit;
                    pathState1.unitPath.add(pathState1.node);
                    pathStates.add(pathState1);
                }
            } else{
                // 如果是其它语句，那么就是一条路径
                // 生成一个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                // 生成一个新的约束
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);
            }
        }
    }

    private boolean pathCheck(Context context, BranchPathState finalNode) {
        List<BoolExpr> pathConstraint = calPathConstraint(context, finalNode.unitPath);
        Solver solver = context.mkSolver();
        for (BoolExpr expr : pathConstraint) {
            solver.add(expr);
        }
        if (solver.check() == Status.UNSATISFIABLE) return false;

        for(Pair<Unit,Unit> pair : finalNode.branchRecord){
            if(!visitedIfNodeNext.contains(pair)){
                return true;
            }
        }
        return false;
    }
}
